from typing import Literal
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from patagent.llm import chat_llm
from patagent.prompt import ROUTER_PROMPT

# Data model
class RouteQuery(BaseModel):
    """Route a user query to the most relevant datasource."""

    datasource: Literal["vectorstore", "chitchat", "data"] = Field(
        ...,
        description="Given a user question choose to route it to chitchat or a vectorstore or data request.",
    )


# LLM with function call
structured_llm_router = chat_llm.with_structured_output(RouteQuery)

route_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", ROUTER_PROMPT),
        ("human", "{question}"),
    ]
)

question_router = route_prompt | structured_llm_router